#------------------------------------------------------------------------------
#------------------------------------------------------------------------------
#CERTAINTY EQUIVALENT CONSUMPTION - COMPARISON ACROSS MODELS
#------------------------------------------------------------------------------
#------------------------------------------------------------------------------

#--------------------------------------------------------------------------------
function cec_fx(ce, W_ctfl::Array{Float64,2}, W_inf_friction::Array{Float64,3})
     cc     = ( repeat(W_ctfl, outer=[1,1, ce.nz]) ./ W_inf_friction  ).^(1/(1-ce.γ)) - 1
     return cc
end
#--------------------------------------------------------------------------------


using JLD
using Interpolations
using PyPlot

#LOAD BASELINE ECONOMY
#--------------------------------------------------------------------------------
name_file = string("../01_Baseline_Model/model_data/CE_model_structure_alpha_",0.028,".jld");
ce        = load(string(name_file), "ce");

#LOAD PERFECT INFORMATION & CONSTANT TYPES CASES
#--------------------------------------------------------------------------------
ce_constant     = load("../02_Perfect_Information/model_data/CE_model_ConsTypes.jld", "ce")
WP_constant, WI_constant = ce_constant.Wf[:,:,1], ce_constant.Wf[:,:,2]
knots           = (ce_constant.b_grid, ce_constant.y_grid)
FF_WP_constant  = interpolate(knots, ce_constant.Wf[:,:,1],      Gridded(Linear()));
FF_WI_constant  = interpolate(knots, ce_constant.Wf[:,:,2],      Gridded(Linear()));
WP_constant     = FF_WP_constant[ce.b_grid,ce.y_grid]
WI_constant     = FF_WI_constant[ce.b_grid,ce.y_grid]

ce_perfinfo     = load("../02_Perfect_Information/model_data/CE_model_PerfInfo.jld", "ce")
WP_perfinfo, WI_perfinfo = ce_perfinfo.Wf[:,:,1], ce_perfinfo.Wf[:,:,2]
knots           = (ce_perfinfo.b_grid, ce_perfinfo.y_grid)
FF_WP_perfinfo  = interpolate(knots, ce_perfinfo.Wf[:,:,1],      Gridded(Linear()));
FF_WI_perfinfo  = interpolate(knots, ce_perfinfo.Wf[:,:,2],      Gridded(Linear()));
WP_perfinfo     = FF_WP_perfinfo[ce.b_grid,ce.y_grid]
WI_perfinfo     = FF_WI_perfinfo[ce.b_grid,ce.y_grid]


#Choose values for plots
i_z_H      = searchsortedfirst(ce.z_grid,0.80)
i_z_L      = searchsortedfirst(ce.z_grid,0.20)
i_b_L      = searchsortedfirst(ce.b_grid,0.45)
i_b_H      = searchsortedfirst(ce.b_grid,0.80)
i_y_avg    = searchsortedfirst(ce.y_grid,1.00)

#--------------------------------------------------------------------------------------------------
#Comparison with the perfect information case
#--------------------------------------------------------------------------------------------------
#Note 1: If w>0, the goverment prefers the perfect information case
#Note 2: Ex-ante expected utility; before the type {S,C} is drawn from nature
#--------------------------------------------------------------------------------------------------
figure()
cc_PerfectInfo  = cec_fx(ce, 0.5*WP_perfinfo+0.5*WI_perfinfo,   0.5*ce.WP+0.5*ce.WI)
PyPlot.figure()
PyPlot.surf(ce.z_grid, ce.b_grid[i_b_L:i_b_H], 100*cc_PerfectInfo[i_b_L:i_b_H,i_y_avg,:], cmap=ColorMap("jet"), alpha=0.7)
xlabel("ζ")
ylabel("b")
zlabel("CEC, %")
tight_layout()
PyPlot.savefig("figures/CEC_Baseline_vs_PerfectInfo.pdf")
#--------------------------------------------------------------------------------------------------


#--------------------------------------------------------------------------------------------------
#Comparison with the constant type case
#--------------------------------------------------------------------------------------------------
#Note: If w>0, the j-type goverment prefers the perfect information case
#--------------------------------------------------------------------------------------------------
#C-type
#-------
     cc_constant_P     = cec_fx(ce, WP_constant,   ce.WP)
     PyPlot.figure()
     PyPlot.surf(ce.z_grid, ce.b_grid[i_b_L:i_b_H], 100*cc_constant_P[i_b_L:i_b_H,i_y_avg,:], cmap=ColorMap("jet"), alpha=0.7)
     xlabel("ζ")
     ylabel("b")
     zlabel("CEC, %")
     tight_layout()
     PyPlot.savefig("figures/CEC_Baseline_vs_ConstantTypes_P.pdf")
#S-type
#-------
     cc_constant_I     = cec_fx(ce, WI_constant,   ce.WI)
     PyPlot.figure()
     PyPlot.surf(ce.z_grid, ce.b_grid[i_b_L:i_b_H], 100*cc_constant_I[i_b_L:i_b_H,i_y_avg,:], cmap=ColorMap("jet"), alpha=0.7)
     xlabel("ζ")
     ylabel("b")
     zlabel("CEC, %")
     tight_layout()
     PyPlot.savefig("figures/CEC_Baseline_vs_ConstantTypes_I.pdf")
#--------------------------------------------------------------------------------------------------
